"""
PostgreSQL数据库备份模块

提供PostgreSQL数据库的备份功能，支持数据导出和结构备份。
"""

import os
import subprocess
import tempfile
import psycopg2
from typing import List, Dict, Any, Optional
from ..core import DatabaseBackup


class PostgreSQLBackup(DatabaseBackup):
    """PostgreSQL数据库备份类
    
    实现PostgreSQL数据库的备份功能。
    """
    
    def __init__(self, host: str, port: int, username: str, password: str, 
                 database: str, backup_dir: str = "backups",
                 pg_dump_path: str = "pg_dump"):
        """初始化PostgreSQL备份实例
        
        参数:
            host (str): PostgreSQL主机地址
                示例值: "localhost"
            port (int): PostgreSQL端口
                示例值: 5432
            username (str): PostgreSQL用户名
                示例值: "postgres"
            password (str): PostgreSQL密码
                示例值: "password123"
            database (str): 数据库名称
                示例值: "my_database"
            backup_dir (str): 备份文件存储目录
                示例值: "backups"
            pg_dump_path (str): pg_dump命令路径
                示例值: "pg_dump"
        
        返回:
            PostgreSQLBackup: PostgreSQL备份实例
        """
        super().__init__(host, port, username, password, database, backup_dir)
        self.pg_dump_path = pg_dump_path
        self.connection = None
    
    def connect(self) -> bool:
        """连接到PostgreSQL数据库
        
        返回:
            bool: 连接是否成功
                示例值: True
        """
        try:
            self.connection = psycopg2.connect(
                host=self.host,
                port=self.port,
                user=self.username,
                password=self.password,
                database=self.database
            )
            self.logger.info(f"成功连接到PostgreSQL数据库 {self.database}")
            return True
        except Exception as e:
            self.logger.error(f"连接PostgreSQL数据库失败: {e}")
            return False
    
    def disconnect(self) -> bool:
        """断开PostgreSQL数据库连接
        
        返回:
            bool: 断开是否成功
                示例值: True
        """
        try:
            if self.connection:
                self.connection.close()
                self.connection = None
                self.logger.info("已断开PostgreSQL数据库连接")
            return True
        except Exception as e:
            self.logger.error(f"断开PostgreSQL数据库连接失败: {e}")
            return False
    
    def get_tables(self) -> List[str]:
        """获取PostgreSQL数据库中的所有表名
        
        返回:
            List[str]: 表名列表
                示例值: ["users", "products", "orders"]
        """
        if not self.connection:
            self.logger.error("数据库未连接")
            return []
        
        try:
            with self.connection.cursor() as cursor:
                cursor.execute("""
                    SELECT table_name 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public' 
                    AND table_type = 'BASE TABLE'
                    ORDER BY table_name
                """)
                tables = [row[0] for row in cursor.fetchall()]
                self.logger.info(f"获取到 {len(tables)} 个表")
                return tables
        except Exception as e:
            self.logger.error(f"获取表列表失败: {e}")
            return []
    
    def backup_table(self, table_name: str, output_file: str, 
                    include_data: bool = True) -> bool:
        """备份单个PostgreSQL数据表
        
        参数:
            table_name (str): 表名
                示例值: "users"
            output_file (str): 输出文件路径
                示例值: "backups/users_20231201.sql"
            include_data (bool): 是否包含数据
                示例值: True
        
        返回:
            bool: 备份是否成功
                示例值: True
        """
        try:
            # 构建pg_dump命令
            cmd = [
                self.pg_dump_path,
                f"-h{self.host}",
                f"-p{self.port}",
                f"-U{self.username}",
                f"-d{self.database}",
                f"-t{table_name}",
                "--no-owner",
                "--no-privileges"
            ]
            
            # 如果不包含数据，只备份结构
            if not include_data:
                cmd.append("--schema-only")
            
            # 设置环境变量（用于密码）
            env = os.environ.copy()
            env['PGPASSWORD'] = self.password
            
            # 执行备份命令
            with open(output_file, 'w', encoding='utf-8') as f:
                result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, 
                                      env=env, text=True, timeout=300)
            
            if result.returncode == 0:
                self.logger.info(f"成功备份表 {table_name} 到 {output_file}")
                return True
            else:
                self.logger.error(f"备份表 {table_name} 失败: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            self.logger.error(f"备份表 {table_name} 超时")
            return False
        except Exception as e:
            self.logger.error(f"备份表 {table_name} 时发生错误: {e}")
            return False
    
    def backup_database(self, output_file: str, include_data: bool = True,
                       tables: Optional[List[str]] = None) -> Dict[str, Any]:
        """备份整个PostgreSQL数据库
        
        参数:
            output_file (str): 输出文件路径
                示例值: "backups/full_backup_20231201.sql"
            include_data (bool): 是否包含数据
                示例值: True
            tables (Optional[List[str]]): 指定要备份的表（None表示备份所有表）
                示例值: ["users", "products"]
        
        返回:
            Dict[str, Any]: 备份结果信息
                示例值: {
                    "success": True,
                    "output_file": "backups/full_backup_20231201.sql",
                    "tables_backed_up": 10,
                    "file_size": "15.2MB"
                }
        """
        try:
            # 构建pg_dump命令
            cmd = [
                self.pg_dump_path,
                f"-h{self.host}",
                f"-p{self.port}",
                f"-U{self.username}",
                f"-d{self.database}",
                "--no-owner",
                "--no-privileges",
                "--verbose"
            ]
            
            # 如果不包含数据，只备份结构
            if not include_data:
                cmd.append("--schema-only")
            
            # 如果指定了表，只备份这些表
            if tables:
                for table in tables:
                    cmd.extend(["-t", table])
            
            # 设置环境变量（用于密码）
            env = os.environ.copy()
            env['PGPASSWORD'] = self.password
            
            # 执行备份命令
            with open(output_file, 'w', encoding='utf-8') as f:
                result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, 
                                      env=env, text=True, timeout=1800)  # 30分钟超时
            
            if result.returncode == 0:
                file_size = os.path.getsize(output_file)
                file_size_str = self._format_file_size(file_size)
                
                result_info = {
                    "success": True,
                    "output_file": output_file,
                    "file_size": file_size_str,
                    "tables_backed_up": len(tables) if tables else "all"
                }
                
                self.logger.info(f"成功备份数据库 {self.database} 到 {output_file}")
                return result_info
            else:
                self.logger.error(f"备份数据库失败: {result.stderr}")
                return {
                    "success": False,
                    "error": result.stderr
                }
                
        except Exception as e:
            # 检查是否是超时异常
            if "timeout" in str(e).lower():
                self.logger.error("备份数据库超时")
                return {
                    "success": False,
                    "error": "备份操作超时"
                }
            else:
                self.logger.error(f"备份数据库时发生错误: {e}")
                return {
                    "success": False,
                    "error": str(e)
                }
    
    def test_connection(self) -> Dict[str, Any]:
        """测试数据库连接
        
        返回:
            Dict[str, Any]: 连接测试结果
                示例值: {
                    "success": True,
                    "database_version": "15.3",
                    "connection_info": {
                        "host": "localhost",
                        "port": 5432,
                        "database": "my_database"
                    }
                }
        """
        try:
            if not self.connect():
                return {
                    "success": False,
                    "error": "连接失败"
                }
            
            with self.connection.cursor() as cursor:
                cursor.execute("SELECT version()")
                version = cursor.fetchone()[0]
                
                cursor.execute("SELECT current_database()")
                current_db = cursor.fetchone()[0]
            
            result = {
                "success": True,
                "database_version": version,
                "current_database": current_db,
                "connection_info": {
                    "host": self.host,
                    "port": self.port,
                    "database": self.database
                }
            }
            
            self.disconnect()
            return result
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def get_database_info(self) -> Dict[str, Any]:
        """获取数据库详细信息
        
        返回:
            Dict[str, Any]: 数据库信息
                示例值: {
                    "tables_count": 15,
                    "total_size": "256MB",
                    "encoding": "UTF8",
                    "collation": "en_US.UTF-8"
                }
        """
        if not self.connect():
            return {
                "tables_count": 0,
                "total_size": "0B",
                "error": "连接失败"
            }
        
        try:
            with self.connection.cursor() as cursor:
                # 获取表数量
                cursor.execute("""
                    SELECT COUNT(*) 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public' 
                    AND table_type = 'BASE TABLE'
                """)
                tables_count = cursor.fetchone()[0]
                
                # 获取数据库大小
                cursor.execute("SELECT pg_database_size(%s)", (self.database,))
                total_size_bytes = cursor.fetchone()[0]
                
                # 获取编码和排序规则
                cursor.execute("""
                    SELECT encoding, datcollate 
                    FROM pg_database 
                    WHERE datname = %s
                """, (self.database,))
                encoding_info = cursor.fetchone()
                
            result = {
                "tables_count": tables_count,
                "total_size": self._format_file_size(total_size_bytes),
                "encoding": encoding_info[0] if encoding_info else "unknown",
                "collation": encoding_info[1] if encoding_info else "unknown"
            }
            
            self.disconnect()
            return result
            
        except Exception as e:
            self.logger.error(f"获取数据库信息失败: {e}")
            return {
                "tables_count": 0,
                "total_size": "0B",
                "error": str(e)
            }
    
    def backup_schema_only(self, output_file: str) -> Dict[str, Any]:
        """只备份数据库结构（不包含数据）
        
        参数:
            output_file (str): 输出文件路径
                示例值: "backups/schema_only_20231201.sql"
        
        返回:
            Dict[str, Any]: 备份结果信息
                示例值: {
                    "success": True,
                    "output_file": "backups/schema_only_20231201.sql",
                    "file_size": "45.2KB"
                }
        """
        return self.backup_database(output_file, include_data=False)
    
    def list_schemas(self) -> List[str]:
        """获取数据库中的所有模式
        
        返回:
            List[str]: 模式名列表
                示例值: ["public", "admin", "reports"]
        """
        if not self.connect():
            return []
        
        try:
            with self.connection.cursor() as cursor:
                cursor.execute("""
                    SELECT schema_name 
                    FROM information_schema.schemata 
                    WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast')
                    ORDER BY schema_name
                """)
                schemas = [row[0] for row in cursor.fetchall()]
                return schemas
        except Exception as e:
            self.logger.error(f"获取模式列表失败: {e}")
            return []
        finally:
            self.disconnect()